import torch
from torch._C import device
from torch_sparse import SparseTensor as st, diag
import torch_sparse 
from torch_sparse.matmul import matmul as spspmm
from fast_inverse import fast_inverse

class JacobPreconditioner:
    
    def __init__(self,K_indices,jacobi_indices,values,m) -> None:
        
        # A = torch.empty(m,9,dtype=torch.float,device=device)
        A = values[jacobi_indices].contiguous().view(m,9)
        A_inv = fast_inverse(A)
        # A_inv_ = coalesce()
        self.__diag = st(row = K_indices[0,jacobi_indices].contiguous().view(-1), col = K_indices[1,jacobi_indices].contiguous().view(-1),value = A_inv.contiguous().view(-1), sparse_sizes=(m*3,m*3))
        # A=st(row=indices[0],col=indices[1],value=values,sparse_sizes=(m,m),is_sorted=True).to_symmetric()

        # idx=torch.arange(m,dtype=indices.dtype,device=indices.device)
        # self.__diag=st(idx,None,idx,1/diag.get_diag(A))

    def apply(self,indices,values,m,n):
        B=st(row=indices[0],col=indices[1],value=values,sparse_sizes=(m,n),is_sorted=True)

        C=spspmm(self.__diag,B)
        row, col, value = C.coo()
        return torch.stack([row, col], dim=0), value


